# -*- coding: utf-8 -*-
"""
Module for handling the computational mesh.

This module defines the data structures for the mesh and includes functionality
to load a mesh from a file using the meshio library.
"""

import meshio
import numpy as np
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class Mesh:
    """
    A class to hold and manage the computational mesh.
    
    This class loads a mesh from a file and organizes its components, such as
    points (nodes), cells (elements), and physical groups (boundaries). It
    also identifies which parts of the mesh are in the inner rigid zone vs.
    the outer deforming zone.
    """
    def __init__(self, file_path):
        """
        Initializes the Mesh object by loading a mesh file.

        Args:
            file_path (str): The path to the mesh file (e.g., .msh from GMSH).
        """
        self.file_path = file_path
        self.points = None
        self.cells = None
        self.cell_data = {}
        self.field_data = {}
        self.gmsh_physical_tags = {}
        
        # Node and cell sets for different mesh zones
        self.rigid_zone_nodes = None
        self.deforming_zone_nodes = None
        self.wing_nodes = None

        try:
            logging.info(f"Loading mesh from {self.file_path}...")
            self.mesh_raw = meshio.read(self.file_path)
            self._process_mesh()
            logging.info("Mesh loaded and processed successfully.")
            self._identify_mesh_zones()
        except FileNotFoundError:
            logging.error(f"Mesh file not found at {self.file_path}.")
            raise
        except Exception as e:
            logging.error(f"An error occurred while loading or processing the mesh: {e}")
            raise

    def _process_mesh(self):
        """
        Extracts relevant information from the raw meshio object.
        """
        self.points = self.mesh_raw.points
        self.cells = self.mesh_raw.cells_dict
        
        # Store physical tags from GMSH if they exist
        if self.mesh_raw.field_data:
            for name, (tag, dim) in self.mesh_raw.field_data.items():
                self.gmsh_physical_tags[name] = int(tag)
        
        logging.info(f"  - Points: {len(self.points)}")
        logging.info(f"  - Cell types: {list(self.cells.keys())}")
        logging.info(f"  - Physical tags found: {self.gmsh_physical_tags}")

    def _identify_mesh_zones(self):
        """
        Identifies nodes belonging to the rigid (inner) and deforming (outer) zones.
        
        This is a placeholder implementation. A real implementation would need to
        rely on physical tags defined in the mesh file during generation (e.g., in GMSH).
        For example, one volume would be tagged as 'inner_zone' and another as 'outer_zone'.
        """
        logging.info("Identifying mesh zones (rigid vs. deforming)...")
        
        # Initialize node arrays
        num_nodes = self.points.shape[0]
        self.deforming_zone_nodes = np.arange(num_nodes)
        self.rigid_zone_nodes = np.array([], dtype=int)
        self.wing_nodes = np.array([], dtype=int)

        # Try to identify wing nodes from physical tags
        wing_surface_tag = self.gmsh_physical_tags.get('wing_surface')
        if wing_surface_tag:
            logging.info(f"Found 'wing_surface' tag with ID {wing_surface_tag}")
            try:
                # Look for cells with the wing surface tag
                wing_node_set = set()
                
                # Check all cell blocks in the mesh
                for i, cell_block in enumerate(self.mesh_raw.cells):
                    cell_type = cell_block.type
                    cells = cell_block.data
                    
                    # Check if this cell block has physical data
                    if hasattr(self.mesh_raw, 'cell_data') and 'gmsh:physical' in self.mesh_raw.cell_data:
                        physical_tags = self.mesh_raw.cell_data['gmsh:physical'][i]
                        # Find cells with wing surface tag
                        wing_cells = np.where(physical_tags == wing_surface_tag)[0]
                        if len(wing_cells) > 0:
                            # Get all nodes from these cells
                            for cell_idx in wing_cells:
                                wing_node_set.update(cells[cell_idx])
                
                self.wing_nodes = np.array(list(wing_node_set), dtype=int)
                logging.info(f"  - Identified {len(self.wing_nodes)} nodes on wing surface")
                
                # For now, treat wing nodes as part of rigid zone
                if len(self.wing_nodes) > 0:
                    self.rigid_zone_nodes = self.wing_nodes
                    # Remove wing nodes from deforming zone
                    self.deforming_zone_nodes = np.setdiff1d(self.deforming_zone_nodes, self.wing_nodes)
                    
            except Exception as e:
                logging.warning(f"Failed to extract wing nodes: {e}")
                logging.warning("No 'wing_surface' physical tag found. Cannot identify wing nodes.")
        else:
            logging.warning("No 'wing_surface' physical tag found. Cannot identify wing nodes.")

        logging.info(f"  - Nodes in rigid zone: {len(self.rigid_zone_nodes)}")
        logging.info(f"  - Nodes in deforming zone: {len(self.deforming_zone_nodes)}")

    def get_boundary_nodes(self, boundary_name: str) -> np.ndarray:
        """
        Gets the node indices for a named physical boundary.

        Args:
            boundary_name (str): The name of the boundary (e.g., "wing_surface", "inlet").

        Returns:
            np.ndarray: An array of node indices belonging to that boundary.
                         Returns an empty array if the boundary is not found.
        """
        boundary_tag = self.gmsh_physical_tags.get(boundary_name)
        if boundary_tag is None:
            logging.warning(f"Boundary with name '{boundary_name}' not found in physical tags.")
            return np.array([], dtype=int)

        boundary_nodes = set()
        # Find the cell block associated with this physical tag
        for cell_type, cell_block in self.cells.items():
            # meshio doesn't directly link physical tags to cell blocks in a simple way.
            # This requires a more complex parsing of the mesh structure, often
            # specific to the file format (like GMSH .msh format).
            # This is a significant challenge and often requires a dedicated library
            # or more complex code.
            pass
        
        if not boundary_nodes:
            logging.warning(f"Could not extract nodes for boundary '{boundary_name}'. "
                            f"This part of the code is a placeholder and needs a robust implementation.")

        return np.array(list(boundary_nodes), dtype=int)


    def write(self, file_path, point_data=None, cell_data=None):
        """
        Writes the current state of the mesh to a file (e.g., VTK for ParaView).

        Args:
            file_path (str): The path to the output file.
            point_data (dict): Dictionary of data to be associated with points.
            cell_data (dict): Dictionary of data to be associated with cells.
        """
        try:
            mesh_to_write = meshio.Mesh(
                self.points,
                self.mesh_raw.cells, # Use original cell connectivity
                point_data=point_data,
                cell_data=cell_data
            )
            logging.debug(f"Writing mesh to {file_path}...")
            meshio.write(file_path, mesh_to_write)
        except Exception as e:
            logging.error(f"Failed to write mesh file to {file_path}: {e}")
            raise

if __name__ == '__main__':
    # Example usage:
    # This requires a dummy mesh file to exist at 'mesh/dummy.msh'
    # You would replace this with your actual mesh file from GMSH.
    import os
    if not os.path.exists("mesh"):
        os.makedirs("mesh")
    try:
        # Create a simple dummy mesh for testing if one doesn't exist
        if not os.path.exists("mesh/dummy.msh"):
            points = [[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]]
            cells = [("triangle", [[0, 1, 2], [0, 2, 3]])]
            meshio.write("mesh/dummy.msh", meshio.Mesh(points=points, cells=cells))
            
        mesh = Mesh("mesh/dummy.msh")
        print("\n--- Mesh Object Details ---")
        print(f"Number of points: {len(mesh.points)}")
        print(f"Cell dictionary: {mesh.cells}")
    except Exception as e:
        print(f"Could not run example. Make sure you have a mesh file available. Error: {e}") 